#include "iv-res.hpp"
#include <tree-pass.h>
#include <print-rtl.h>
#include <print-tree.h>
#include <cfgloop.h>
#include <gimple-ssa.h>
#include <context.h>
#include <tree-core.h>
#include <tree-vrp.h>
#include <stringpool.h>
#include <tree-ssanames.h>
#include <tree-ssa-alias.h>
#include <gimple-pretty-print.h>

#include <set>
namespace {
  const pass_data pass_data_sw_64_iv_alias = {
    RTL_PASS,
    "iv_alias",
    OPTGROUP_LOOP,
    TV_NONE,
    (PROP_cfg | PROP_ssa),
    0,
    0,
    0,
    0
  };

  const pass_data pass_data_sw_64_early_iv_alias = {
    GIMPLE_PASS,
    "early_iv_alias",
    OPTGROUP_LOOP,
    TV_NONE,
    (PROP_cfg | PROP_ssa),
    0,
    0,
    0,
    0
  };

};

class pass_sw_64_iv_alias : public rtl_opt_pass {
public:
  pass_sw_64_iv_alias(gcc::context *ctx) : rtl_opt_pass(pass_data_sw_64_iv_alias, ctx){}
  virtual bool gate(function *){
    return true;
  }
  void adjust(rtx_insn *insn, rtx x, std::set<tree> &local_iv_vars){
    struct rtx_iv iv;
    if (MEM_P(x)) {
      rtx reg = XEXP(x, 0);
      if (GET_CODE(reg) == PLUS){
        reg = XEXP(reg, 0);
      }

      if (GET_CODE(reg) != REG) {
        return;
      }
      
      // debug(x);
      // debug_tree(MEM_EXPR(x));
      // debug(reg);
      
      if (iv_analyze(insn, reg, &iv)){
        rtx base = iv.base;
        while (!REG_P(base) && GET_CODE(base) == PLUS){
          //if (CONST_INT_P(XEXP(base, 1))){
          base = XEXP(base, 0);
            //}
        }
        if (REG_P(base) && local_iv_vars.find(REG_EXPR(base)) != local_iv_vars.end()){
            SSA_NAME_PTR_INFO(TREE_OPERAND(MEM_EXPR(x), 0))->pt.escaped = 0;
            SSA_NAME_PTR_INFO(TREE_OPERAND(MEM_EXPR(x), 0))->pt.nonlocal = 0;
        } else {
          // puts("IV ANALYZED but UNRECOGNIZED");
          // debug(x);
          // debug_tree(MEM_EXPR(x));
          // debug(reg);
          // debug(iv.base);
        }
      } else {
        // puts("IV not ANALYZED");
        // debug(x);
        // debug_tree(MEM_EXPR(x));
        // debug(reg);
        // debug(REG_EXPR(reg));
      }
    } else {
      // puts("not mem:");
      // debug(x);
    }
  }
  unsigned int execute(function *f){
    struct loop* loop;
    std::set<tree> local_var_iv_tree;
    for (tree *t = cfun->gimple_df->ssa_names->begin(); t != cfun->gimple_df->ssa_names->end(); t ++){
      if (*t && SSA_NAME_VAR(*t)){
        tree var = SSA_NAME_VAR(*t);
        if (DECL_NAME(var) && IDENTIFIER_POINTER(DECL_NAME(var)) && !strncmp(IDENTIFIER_POINTER(DECL_NAME(var)), "ivtmp", 5)){

          gimple *def_stmt = SSA_NAME_DEF_STMT(*t);
          //debug(def_stmt);
          if (gimple_code(def_stmt) == GIMPLE_ASSIGN)
            if ((gimple_assign_rhs_code(def_stmt) == NOP_EXPR || gimple_assign_rhs_code(def_stmt) == PLUS_EXPR)){
              auto rhsexpr = gimple_assign_rhs1(def_stmt);

              if (TREE_CODE(rhsexpr) == ADDR_EXPR){
                auto rhsbase = TREE_OPERAND(rhsexpr, 0);
                if (VAR_P(rhsbase)){
                  if (DECL_CONTEXT(rhsbase) == cfun->decl){
                    local_var_iv_tree.insert(var);
                    puts("insert");
                    debug_tree(var);

                  }
                }
              } else if (TREE_CODE(rhsexpr) == SSA_NAME) {
                tree varbase = SSA_NAME_VAR(rhsexpr);
                if (varbase && DECL_NAME(varbase) && IDENTIFIER_POINTER(DECL_NAME(varbase)) && !strncmp(IDENTIFIER_POINTER(DECL_NAME(varbase)), "ivtmp", 5)){
                  puts("insert");
                  debug_tree(varbase);
                  if (local_var_iv_tree.find(varbase) != local_var_iv_tree.end()){
                    local_var_iv_tree.insert(var);
                    puts("IV_RHS");
                    debug_tree(*t);
                    debug_tree(rhsexpr);
                  }
                }
              }
            }
        }
      }
    }
    FOR_EACH_LOOP(loop, 0){
      iv_analysis_loop_init(loop);
    }
    FOR_EACH_LOOP(loop, 0){
      auto bbs = get_loop_body(loop);
      for (int i = 0; i < loop->num_nodes; i ++){
        auto bb = bbs[i];
        for (auto insn = BB_HEAD(bb); insn != BB_END(bb); insn = NEXT_INSN(insn)){
          if (GET_CODE(insn) != INSN) continue;
          auto pat = PATTERN(insn);
          if (GET_CODE(pat) == SET){
            auto dst = XEXP(pat, 0);
            auto src = XEXP(pat, 1);
            adjust(insn, dst, local_var_iv_tree);
            adjust(insn, src, local_var_iv_tree);
          }
        }
      }
    }
    FOR_EACH_LOOP(loop, 0) {
      iv_analysis_done();
    }
    return 0;
  }
};


void sw_64_iv_res::add_passes(){
  //nohp::register_pass(new pass_sw_64_early_iv_alias(g), "ivopts", 1, PASS_POS_INSERT_AFTER);
  //nohp::register_pass(new pass_sw_64_iv_alias(g), "loop2_init", 1, PASS_POS_INSERT_AFTER);
}

class pass_sw_64_early_iv_alias : public gimple_opt_pass {
public:
  pass_sw_64_early_iv_alias(gcc::context *ctx) : gimple_opt_pass(pass_data_sw_64_early_iv_alias, ctx){}
  virtual bool gate(function *){
    return true;
  }
  void adjust(rtx_insn *insn, rtx x, std::set<tree> &local_iv_vars){
    struct rtx_iv iv;
    if (MEM_P(x)) {
      rtx reg = XEXP(x, 0);
      if (GET_CODE(reg) == PLUS){
        reg = XEXP(reg, 0);
      }
      if (GET_CODE(reg) != REG) return;
      
      // debug(x);
      // debug_tree(MEM_EXPR(x));
      // debug(reg);
      
      if (iv_analyze(insn, reg, &iv)){
        // puts("base:");
        // debug(iv.base);
        if (REG_P(iv.base)){
          // puts("reg:");
          // debug_tree(REG_EXPR(iv.base));
          if (local_iv_vars.find(REG_EXPR(iv.base)) != local_iv_vars.end()){
            SSA_NAME_PTR_INFO(TREE_OPERAND(MEM_EXPR(x), 0))->pt.escaped = 0;
            SSA_NAME_PTR_INFO(TREE_OPERAND(MEM_EXPR(x), 0))->pt.nonlocal = 0;
          }
        }
      }
    }    
  }
  unsigned int execute(function *f){
    struct loop* loop;
    std::set<tree> local_var_iv_tree;
    for (tree *t = cfun->gimple_df->ssa_names->begin(); t != cfun->gimple_df->ssa_names->end(); t ++){
      if (*t && SSA_NAME_VAR(*t)){
        tree var = SSA_NAME_VAR(*t);
        if (DECL_NAME(var) && IDENTIFIER_POINTER(DECL_NAME(var)) && !strncmp(IDENTIFIER_POINTER(DECL_NAME(var)), "ivtmp", 5)){
          gimple *def_stmt = SSA_NAME_DEF_STMT(*t);
          if (gimple_code(def_stmt) == GIMPLE_ASSIGN)
            if ((gimple_assign_rhs_code(def_stmt) == NOP_EXPR || gimple_assign_rhs_code(def_stmt) == PLUS_EXPR)){
              auto rhsexpr = gimple_assign_rhs1(def_stmt);
              if (TREE_CODE(rhsexpr) == ADDR_EXPR){
                //debug_tree(rhsexpr);
                auto rhsbase = TREE_OPERAND(rhsexpr, 0);
                if (VAR_P(rhsbase)){
                  if (DECL_CONTEXT(rhsbase) == cfun->decl){
                    local_var_iv_tree.insert(var);
                  }
                }
              }
            }
        }
      }
    }
    for (auto si = gsi_start(f->gimple_body); !gsi_end_p(si); gsi_next(&si)){
      debug(gsi_stmt(si));
    }
    return 0;
  }
  //   FOR_EACH_LOOP(loop, 0){
  //     iv_analysis_loop_init(loop);
  //     auto bbs = get_loop_body(loop);
  //     for (int i = 0; i < loop->num_nodes; i ++){
  //       auto bb = bbs[i];
  //       for (auto insn = BB_HEAD(bb); insn != BB_END(bb); insn = NEXT_INSN(insn)){
  //         if (GET_CODE(insn) != INSN) continue;
  //         auto pat = PATTERN(insn);
  //         if (GET_CODE(pat) == SET){
  //           auto dst = XEXP(pat, 0);
  //           auto src = XEXP(pat, 1);
  //           adjust(insn, dst, local_var_iv_tree);
  //           adjust(insn, src, local_var_iv_tree);
  //         }
  //       }
  //     }

  //   }
  //   return 0;
  // }
};
